#pragma once

#include <Python.h>
#include <ATen/Operators.h>
#include <torch/library.h>

// 将两个变量的名字拼接
#define _CONCAT(A, B) A##B                   //! "##"用于将两个标记连接在一起
//! 当AB本身也是宏时，_CONCAT会直接将宏名拼接，而不会进行展开因此需要定义CONCAT
//! CONCAT会优先将AB展开，再进行拼接
//! 这种方法我们称为宏的分层定义
#define CONCAT(A, B) _CONCAT(A, B)

#define _STRINGIFY(A) #A                     //! "#"用于将变量字符串化
#define STRINGIFY(A) _STRINGIFY(A)

//! 这里也通用宏分层定义，前期采用TORCH_LIBRARY(TORCH_EXTENSION_NAME, m)的形式踩过坑
// TORCH_LIBRARY(namespace, torch::Library)
// 决定了算子所正在的空间，Python调用时必须采用torch.ops.namespace.func_name的方式
#define TORCH_LIBRARY_EXPAND(NAME, MODULE)                \
    TORCH_LIBRARY(NAME, MODULE)              // 向PyTorch声明NAME命名空间下的函数签名


#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE)   \
    TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) // 向PyTorch注册NAME命名空间下的具体函数实现


#define REGISTER_EXTENSION(NAME)                                     \
    PyMODINIT_FUNC CONCAT(PyInit_, NAME)(void) {                     \
        static struct PyModuleDef module = {PyModuleDef_HEAD_INIT,   \
                                        STRINGIFY(NAME),             \
                                        nullptr,                     \
                                        0,                           \
                                        nullptr};                    \
        return PyModule_Create(&module);                             \
    }